import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Bidirectional, LSTM, Dense, Dropout
from tensorflow.keras.callbacks import EarlyStopping, ReduceLROnPlateau
from tensorflow.keras.losses import Huber
from sklearn.metrics import mean_absolute_error, mean_squared_error
import numpy as np
import plotly.graph_objects as go
import pandas as pd
from data_pipeline import DataPipeline


class LSTMModel:
    def __init__(self, input_shape):
        self.model = self._build_model(input_shape)

    def _build_model(self, input_shape):
        model = Sequential([
            Bidirectional(LSTM(128, return_sequences=True), input_shape=input_shape),
            Dropout(0.3),
            LSTM(64, return_sequences=True),
            Dropout(0.2),
            LSTM(32),
            Dense(32, activation='relu'),
            Dense(5)
        ])

        model.compile(
            optimizer=tf.keras.optimizers.Adam(learning_rate=0.001),
            loss=Huber(),
            metrics=['mae']
        )
        return model

    def train(self, X_train, y_train, X_val, y_val):
        """模型训练"""
        callbacks = [
            EarlyStopping(patience=25, restore_best_weights=True, monitor='val_loss'),
            ReduceLROnPlateau(factor=0.5, patience=5, verbose=1)
        ]

        history = self.model.fit(
            X_train, y_train,
            validation_data=(X_val, y_val),
            epochs=100,
            batch_size=32,
            callbacks=callbacks,
            verbose=1
        )
        return history

    def predict(self, X):
        return self.model.predict(X)

class ForecastingSystem:
    def __init__(self, data_pipeline=None):
        self.data_pipeline = data_pipeline if data_pipeline else DataPipeline()
        self.full_data = None
        self.scaler_pipeline = None
        self.model = None

    def _create_sequences(self, data):
        """创建时间序列数据集"""
        X, y = [], []
        for i in range(len(data) - 30 - 5):
            X.append(data[i:i + 30])
            y.append(data[i + 30:i + 30 + 5, 3])
        return np.array(X), np.array(y)

    def prepare_data(self):
        # 获取处理后的数据
        df = self.data_pipeline.get_processed_data()
        self.full_data = df

        # 数据集划分
        train_data, test_data = train_test_split(df, test_size=0.15, shuffle=False)

        # 初始化标准化器
        self.scaler_pipeline = ScalerPipeline(train_data)

        # 转换数据
        scaled_train = self.scaler_pipeline.transform(train_data)
        scaled_test = self.scaler_pipeline.transform(test_data)

        # 创建序列
        X_train, y_train = self._create_sequences(scaled_train)
        X_test, y_test = self._create_sequences(scaled_test)

        # 验证集划分
        X_train, X_val, y_train, y_val = train_test_split(X_train, y_train, test_size=0.2, shuffle=False)

        return (X_train, y_train), (X_val, y_val), (X_test, y_test)

    def prepare_data_from_request(self, data):
        # 假设data是一个包含股票数据的字典
        # 这里需要根据实际数据格式进行处理
        df = pd.DataFrame(data)
        # ... 数据预处理代码 ...
        return (X_train, y_train), (X_val, y_val), (X_test, y_test)

    def train_model(self, X_train, y_train, X_val, y_val):
        """模型训练"""
        self.model = LSTMModel((30, X_train.shape[2]))
        history = self.model.train(X_train, y_train, X_val, y_val)
        return history

    def _inverse_transform(self, scaled_values):
        """逆变换价格数据"""
        price_scaler = self.scaler_pipeline.scalers['price']
        close_scale = price_scaler.scale_[3]
        close_center = price_scaler.center_[3]
        return (scaled_values * close_scale) + close_center

    def evaluate(self, X_test, y_test):
        """模型评估"""
        y_pred = self.model.predict(X_test)

        # 逆变换
        y_pred_original = self._inverse_transform(y_pred)
        y_true_original = self._inverse_transform(y_test)

        # 计算指标
        mae = mean_absolute_error(y_true_original, y_pred_original)
        rmse = np.sqrt(mean_squared_error(y_true_original, y_pred_original))

        # 方向准确性
        directional_acc = np.mean(
            np.sign(y_true_original[:, -1] - y_true_original[:, 0]) ==
            np.sign(y_pred_original[:, -1] - y_pred_original[:, 0])
        )

        print(f"\n评估结果:")
        print(f"MAE: {mae:.4f}")
        print(f"RMSE: {rmse:.4f}")
        print(f"方向准确性: {directional_acc:.2%}")

        # 可视化最近50个预测
        self._plot_predictions(y_true_original[-50:], y_pred_original[-50:])

    def _plot_predictions(self, true, pred):
        fig = go.Figure()

        # 真实值
        fig.add_trace(go.Scatter(
            x=np.arange(len(true)),
            y=true[:, -1],
            name='True Price',
            line=dict(color='blue')
        ))

        # 预测值
        fig.add_trace(go.Scatter(
            x=np.arange(len(pred)),
            y=pred[:, -1],
            name='Predicted Price',
            line=dict(color='red', dash='dot')
        ))

        fig.update_layout(
            title='实际价格 vs 预测价格',
            xaxis_title='时间步',
            yaxis_title='价格',
            hovermode='x unified',
            template='plotly_white'
        )
        fig.show()

    def forecast(self, days=5):
        """滚动预测未来价格"""
        if self.model is None:
            raise ValueError("请先训练模型")

        # 获取最新序列
        latest_data = self.full_data[-30:]
        scaled_data = self.scaler_pipeline.transform(latest_data)

        predictions = []
        current_sequence = scaled_data.copy()

        for _ in range(days):
            # 生成预测
            pred = self.model.predict(current_sequence.reshape(1, 30, -1))[0]
            next_close = pred[-1]  # 取最后一个预测值

            # 构建新数据行（保持其他特征合理变化）
            new_row = np.roll(current_sequence, -1, axis=0)

            # 更新Close价格
            new_row[-1, 3] = next_close

            # 模拟其他特征的合理变化（示例逻辑，可根据需求调整）
            new_row[-1, 0] = new_row[-2, 3]  # Open ≈ 前一日Close
            new_row[-1, 1] = new_row[-2, 3] * 1.01  # High假设波动1%
            new_row[-1, 2] = new_row[-2, 3] * 0.99  # Low假设波动1%

            current_sequence = new_row
            predictions.append(next_close)

        # 逆变换预测结果
        predicted_prices = self._inverse_transform(np.array(predictions))

        # 生成预测日期
        last_date = self.full_data.index[-1]
        dates = pd.date_range(start=last_date + pd.Timedelta(days=1), periods=days)

        return pd.Series(predicted_prices, index=dates, name='Predicted Close')